import torch
from typing import Dict, List, Optional, Union

from mmcv.transforms import BaseTransform
from mmaction.registry import TRANSFORMS

try:
    import pytorch_wavelets
except ImportError:
    raise ImportError(
        'Please run "pip install pywavelets pytorch-wavelets" to install Wavelet first.')
from pytorch_wavelets import DWTForward, DWTInverse


@TRANSFORMS.register_module()
class WaveletFilter(BaseTransform):
    """Using wavelet to filter the imags.

    wavelet: https://github.com/dmlc/wavelet

    Required Keys:
        - imgs
        - frame_inds

    Added Keys:
        - imgs
        - original_shape
        - img_shape

    """

    def __init__(self,wavelet='db4', 
                 levels=3, 
                 device='cuda',
                 **kwargs) -> None:
        """
        初始化视频小波处理器
        
        参数:
            wavelet: 小波基类型，默认'db4'
            levels: 小波分解层数, 默认3
            device: 计算设备，默认'cuda'
        """
        self.wavelet = wavelet
        self.levels = levels
        self.device = device
        self.kwargs = kwargs
        # 初始化小波变换和逆变换
        self.dwt = DWTForward(J=self.levels, wave=self.wavelet, mode='zero').to(device)
        self.idwt = DWTInverse(wave=self.wavelet, mode='zero').to(device)
        if self.dwt is None or self.idwt is None:
            raise ValueError('Wavelet filters are not initialized. '
                             'Please set "wavelet_init" in the pipeline.')

    def _filter_coefficients(self, coeffs):
        """
        对小波系数进行滤波处理（优化版本）
        """
        LL = coeffs[0]
        band_coeffs = coeffs[1]
        
        # 对低频分量进行轻微增强
        LL = LL * 1.1
        
        # 对高频分量进行软阈值处理（向量化）
        threshold = 0.05
        for i in range(len(band_coeffs)):
            # 确保当前频带的所有方向系数都是张量
            if not all(isinstance(coeff, torch.Tensor) for coeff in band_coeffs[i]):
                continue
                
            # 一次性处理所有方向的高频系数
            # 使用 tuple() 确保传递给 stack 的是元组
            coeff_stack = torch.stack(tuple(band_coeffs[i]), dim=0)
            sign = torch.sign(coeff_stack)
            abs_val = torch.abs(coeff_stack)
            thresholded = sign * torch.relu(abs_val - threshold)
            
            # 将结果存回列表
            for j in range(len(band_coeffs[i])):
                band_coeffs[i][j] = thresholded[j]
        
        return (LL, band_coeffs)

    def transform(self, results: Dict) -> Dict:
        """Perform the wavelet decoding.

        Args:
            results (dict): The result dict.

        Returns:
            dict: The result dict.
        """
        device = torch.device(self.device if torch.cuda.is_available() else 'cpu')
        device = self.device if self.device == 'cpu' else device

        imgs = results['inputs']
        # 从imgs中获取device
        if not isinstance(imgs, torch.Tensor):
            imgs = torch.tensor(imgs, dtype=torch.float32)

        # 如果imgs不在GPU，将数据移到GPU
        # RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method
        #if imgs.device != 'cuda' and device.type == 'cuda':
        imgs = imgs.to(device='cuda')

        # 确保数据在0-1范围内
        if imgs.max() > 1.0:
            imgs = imgs / 255.0
        
        time_steps, channels, batch_size, height, width = imgs.shape
        
        # 重塑为 [batch_size * time_steps, channels, height, width] 以适应小波变换
        imgs = imgs.reshape(-1, channels, height, width)
        

        # 应用小波变换
        imgs = self.dwt(imgs)
        
        # 对小波系数进行滤波处理
        imgs = self._filter_coefficients(imgs)
        
        # 逆小波变换重构图像
        imgs = self.idwt(imgs)
        
        # 确保数据在正确范围内
        imgs = torch.clamp(imgs, 0, 1)
        
        # 恢复原始形状 [batch_size, time_steps, channels, height, width]
        imgs = imgs.reshape(time_steps, channels, batch_size, height, width)

        imgs = imgs * 255.0
        imgs = imgs.type(torch.uint8)

        results['inputs'] = imgs
        results['original_shape'] = imgs[0].shape[:2]
        results['img_shape'] = imgs[0].shape[:2]

        return results
